import edu.umass.cs.mallet.base.types.*;
import edu.umass.cs.mallet.base.fst.*;
import edu.umass.cs.mallet.base.minimize.*;
import edu.umass.cs.mallet.base.minimize.tests.*;
import edu.umass.cs.mallet.base.pipe.*;
import edu.umass.cs.mallet.base.pipe.iterator.*;
import edu.umass.cs.mallet.base.pipe.tsf.*;
import etxt2db.mallet.MedlineSentence2TokenSequence;
import junit.framework.*;
import java.util.Iterator;
import java.util.Random;
import java.util.regex.*;
import java.io.*;

public class TestGeneTagger {
    int numEvaluations = 0;
    static int iterationsBetweenEvals = 20;
    
    private static String CAPS = "[A-Z]";
    private static String LOW = "[a-z]";
    private static String CAPSNUM = "[A-Z0-9]";
    private static String ALPHA = "[A-Za-z]";
    private static String ALPHANUM = "[A-Za-z0-9]";
    private static String PUNT = "[,\\.;:?!()]";
    private static String QUOTE = "[\"`']";
    
    public static void main (String[] args) throws FileNotFoundException, Exception
    {
	Pipe p = new SerialPipes (new Pipe[] {
	    new MedlineSentence2TokenSequence (),
	    new RegexMatches ("INITCAP", Pattern.compile (CAPS+".*")),
	    new RegexMatches ("CAPITALIZED", Pattern.compile (CAPS+LOW+"*")),
	    new RegexMatches ("ALLCAPS", Pattern.compile (CAPS+"+")),
	    new RegexMatches ("MIXEDCAPS", Pattern.compile ("[A-Z][a-z]+[A-Z][A-Za-z]*")),
	    new RegexMatches ("CONTAINSDIGITS", Pattern.compile (".*[0-9].*")),
	    new RegexMatches ("ALLDIGITS", Pattern.compile ("[0-9]+")),
	    new RegexMatches ("NUMERICAL", Pattern.compile ("[-0-9]+[\\.,]+[0-9\\.,]+")),
	    new RegexMatches ("ALPHNUMERIC", Pattern.compile ("[A-Za-z0-9]+")),
	    new RegexMatches ("CONTAINSDASH", Pattern.compile (ALPHANUM+"+-"+ALPHANUM+"*")),
	    new RegexMatches ("ACRO", Pattern.compile ("[A-Z][A-Z\\.]*\\.[A-Z\\.]*")),
	    new RegexMatches ("LONELYINITIAL", Pattern.compile (CAPS+"\\.")),
	    new RegexMatches ("SINGLECHAR", Pattern.compile (ALPHA)),
	    new RegexMatches ("CAPLETTER", Pattern.compile ("[A-Z]")),
	    new RegexMatches ("PUNC", Pattern.compile (PUNT)),
	    new TokenText ("WORD="),
	    //new LexiconMembership(new File("/home/ryantm/progs/state_list")),
	    new TokenTextCharSuffix("4SUFFIX=",4),
	    new TokenTextCharSuffix("3SUFFIX=",3),
	    new TokenTextCharSuffix("2SUFFIX=",2),
	    new TokenTextCharNGrams ("CHARNGRAM=", new int[] {2,3}),
	    //new FeaturesInWindow("WORD=",-1,1),
	    new FeaturesInWindow("WINDOW=",-1,1,Pattern.compile("WORD=.*"),true), 
	    //new PrintTokenSequenceFeatures(),
	    new TokenSequence2FeatureVectorSequence (true, true)
		});

	if (args[0].equals("train")) {
	    InstanceList trainingData = new InstanceList(p);
	    trainingData.add (new LineGroupIterator (new FileReader (new File (args[1])), Pattern.compile("\n"), true));
	    System.out.println ("Number of predicates in training data: "+p.getDataAlphabet().size());
	    
	    int j = 0;
	    for(;j<trainingData.size();j++){
	    	System.out.println(((Instance)trainingData.get(j)).getData());
	    }
	    // Print out all the target names
	    
	    Alphabet targets = p.getTargetAlphabet();
	    System.out.print ("State labels:");
	    for (int i = 0; i < targets.size(); i++)
		System.out.print (" " + targets.lookupObject(i));
	    System.out.println ("");
	    
	    // Print out some feature information
	    System.out.println ("Number of features = "+p.getDataAlphabet().size());
	    
	    MultiSegmentationEvaluator eval =
		new MultiSegmentationEvaluator (new String[] {"GB"},
						new String[] {"GI"});	
	    
	    //CRFByGISUpdate crf = new CRFByGISUpdate (p, null);
	    CRF3 crf = new CRF3 (p, null);
	    //crf.addFullyConnectedStatesForLabels();
	    //crf.addStatesForLabelsConnectedAsIn (trainingData);
	    crf.addStatesForHalfLabelsConnectedAsIn (trainingData);
	    crf.setGaussianPriorVariance (Double.POSITIVE_INFINITY);
	    //crf.setGaussianPriorVariance (1.0);
	    
	    System.out.println("Training on "+trainingData.size()+" training instances");
	  
	    //crf.train (trainingData, null, null, eval, 350, new AGIS(4.0));
	    crf.train (trainingData, null, null, eval, 350);	    
	    
	    crf.write(new File("model.crf"));
	}
	else if (args[0].equals("test")) {	    
	    // Load the stored CRF model
	    ObjectInputStream ois = 
		new ObjectInputStream(new FileInputStream("model.crf"));
	    //CRFByGISUpdate crf = null;
	    //try { crf = (CRFByGISUpdate)ois.readObject(); }
	    CRF3 crf = null;
	    try { crf = (CRF3) ois.readObject(); }
	    catch (ClassNotFoundException e) { e.printStackTrace(); }
	    ois.close();
	    
	    InstanceList trainingData = new InstanceList (p);
	    trainingData.add(new LineGroupIterator(new FileReader(new File(args[1])), Pattern.compile("^$"), true));
	    
	    InstanceList testData = new InstanceList(p);  
	    testData.add(new LineGroupIterator(new FileReader(new File(args[2])), Pattern.compile("^$"), true));
	    
	    MultiSegmentationEvaluator eval =
		new MultiSegmentationEvaluator (new String[] {"GB"},
						new String[] {"GI"});	

	    eval.evaluate(crf,true,0,true,0.0,null,null,testData);
	    
	    String tagfilename = args[2].substring(0,args[2].indexOf(".list"));

	    System.out.println("Writing tags to "+tagfilename+".crf.tag");
	    
	    FileOutputStream fos = new FileOutputStream(tagfilename+".crf.tag");
	    PrintStream ps = new PrintStream(fos);
	    
	    // Iterate over the sequences in the InstanceList.
	    for (InstanceList.Iterator i = testData.iterator();i.hasNext();) {
	        FeatureVectorSequence sequence = 
	            (FeatureVectorSequence)i.nextInstance().getData();
		
		// Use the CRF to determine most likely labeling of the 
		// sequence.
		Sequence labels = crf.viterbiPath(sequence).output();
		
		// print sequences to file
		for (int j = 0; j < labels.size(); j++) {
		    ps.println(labels.get(j));
	        }
		ps.println();
	    }
	    
	}
	else { System.out.println("bad arguments"); }
    }    
}
